import sqlite3
import zai
from dotenv import load_dotenv
load_dotenv()

sys_prompt="""
            您是 SQL 专家级分析师。在适当的时候，根据用户问题和数据库架构生成 SQL 查询。
            你的输出有且仅能是字符串类型的SQL查询语句
            例如：
            question：介绍一下北京大学
            answer：SELECT * FROM mytable WHERE 中文名字 = '北京大学'      
            数据库架构：
            {}

            数据表记录示例：
            {}          
"""

sys_prompt2="""
    你是一个高考志愿智能问答系统，对于用户提问的问题，你需要按照给出的【查询结果】对问题进行回答。
    你的回答需要按照以下两个步骤：
    1.分析用户问题和查询结果，判断是否有【查询结果】可以解答用户的问题，如果有则说明【查询结果】的名称，如果没有，则首先告知用户没有任何可参考的资料，需要注意答案的准确性。
    2.根据资料内容对问题进行解答，若用户希望根据高考分数得到志愿推荐，那么首先关注学校的投档分，越接近越好，从好的学校推荐，并结合【查询结果】中的学校信息来说明。
    """

user_query2="用户问题:{}\n\n查询结果:\n{}"



def llm_chat(query,sys_prompt,stream=False):
    # llm chat调用
    client = zai.ZaiClient()

    # messages
    messages =[
        {'role':'system','content':sys_prompt},
        {'role':'user','content':query}
    ]

    # chat
    resp = client.chat.completions.create(
        model='glm-4.5-flash',
        messages=messages,
        thinking={
            "type":"disabled", # 关闭深度思考模式
        },
        stream=stream
    )

    # sql指令
    if stream:

        return resp
    else:
        return resp.choices[0].message.content


def execute_sqlschema(db_file,table_name):
    """
    查询sqlite数据库中的表结构
    """
    # 打开数据库文件（加载数据库）
    conn = sqlite3.connect(db_file)
    # conn.row_factory = sqlite3.Row

    # 创建cursor访问表
    cursor = conn.cursor()

    # 查看元数据表
    cursor.execute(f"select sql from sqlite_master where tbl_name='{table_name}'")
    result = cursor.fetchone() # 获取第一个结果

    # 获取表结构
    return result[0]

def execute_sqlcmd(db_file,sql):
    """
    查询sqlite数据库中的表结构
    """
    # 打开数据库文件（加载数据库）
    conn = sqlite3.connect(db_file)
    conn.row_factory = sqlite3.Row

    # 创建cursor访问表
    cursor = conn.cursor()

    # 查看元数据表
    cursor.execute(sql)
    result = cursor.fetchall() # 获取第一个结果

    # # 获取表结构
    # return result

    # 获取JSON格式数据输出
    json_datas=[]
    for row in result:
        record = {}
        for key in row.keys():
            record[key] = row[key]
        json_datas.append(record)
    return json_datas

def query_answer(user_query,tbl_name='data1'):
    # 1.根据数据库和表结构，根据用户问题生成sql指令
    db_file = 'data1.db'

    sql_schema = execute_sqlschema(db_file,tbl_name)
    record_sample = execute_sqlcmd(db_file,f'select * from {tbl_name} limit 1,1')

    sql= llm_chat(user_query,sys_prompt.format(sql_schema,record_sample))


    #2.问题解答提示
    result = execute_sqlcmd(db_file,sql)

    answer = llm_chat(user_query2.format(user_query,result),sys_prompt2,stream=True)

    return answer

if __name__ == '__main__':
    
    # 测试完整查询
    answer = query_answer('了解一下北京航空航天大学')

    for chunk in answer:
        print(chunk.choices[0].delta.content,end='')

